from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.chat_models import ChatTongyi
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

# 设置通义千问API密钥 (从环境变量读取)
AI_QW_KEY = os.getenv("AI_QW_KEY")
AI_QW_URL = os.getenv("AI_QW_URL")

# 设置通义千问的 API Key
os.environ["DASHSCOPE_API_KEY"] = AI_QW_KEY
os.environ["DASHSCOPE_ENDPOINT"] = AI_QW_URL

def query_existing_vector_db(persist_directory: str, query: str) -> dict:
    """
    从已存在的 ChromaDB 中查询数据
    :param persist_directory: ChromaDB 持久化存储路径
    :param query: 用户问题
    :return: 包含查询结果和来源文档的字典
    """
    # 1. 加载已存在的向量数据库
    embeddings = DashScopeEmbeddings(model="text-embedding-v1")
    vector_db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    print("已加载现有的向量数据库")

    # 2. 创建检索式问答链
    llm = ChatTongyi(model="qwen-turbo")

    # 自定义提示模板
    prompt_template = """
    使用以下上下文片段来回答问题。如果你不知道答案，请如实说明。
    上下文：
    {context}

    问题：{question}
    答案：
    """
    PROMPT = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_db.as_retriever(search_kwargs={"k": 10}),  # 检索前3个相关片段
        chain_type_kwargs={"prompt": PROMPT},
        return_source_documents=True
    )

    # 3. 执行查询
    result = qa_chain.invoke({"query": query})
    return result

def main():
    # 使用示例
    persist_directory = "./chroma_db"  # ChromaDB 持久化存储路径
    question = "红薯亩产多少斤"  # 替换为你的查询问题

    result = query_existing_vector_db(persist_directory, question)

    # 打印结果
    print("\n" + "=" * 50)
    print(f"问题: {question}")
    print(f"答案: {result['result']}")
    # print("\n来源文档片段:")
    # for doc in result['source_documents']:
    #     print(f"- 页码 {doc.metadata.get('page', '未知')}: {doc.page_content[:100]}...")

if __name__ == "__main__":
    main()